package com.jd.platform.async.executor;


import com.jd.platform.async.callback.DefaultGroupCallback;
import com.jd.platform.async.callback.IGroupCallback;
import com.jd.platform.async.wrapper.WorkerWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

/**
 * 类入口，可以根据自己情况调整core线程的数量
 * @author wuweifeng wrote on 2019-12-18
 * @version 1.0
 */
public class Async {
    private static final Logger logger = LoggerFactory.getLogger(Async.class);
    /**
     * 默认线程池：有界线程池，防止OOM
     * 核心线程数：CPU核心数 * 2
     * 最大线程数：CPU核心数 * 4
     * 队列容量：1000
     * 拒绝策略：CallerRunsPolicy（让调用者线程执行，提供背压）
     */
    private static final ThreadPoolExecutor COMMON_POOL;
    
    static {
        int cores = Runtime.getRuntime().availableProcessors();
        COMMON_POOL = new ThreadPoolExecutor(
            cores * 2,
            cores * 4,
            60L,
            TimeUnit.SECONDS,
            new java.util.concurrent.LinkedBlockingQueue<>(1000),
            new ThreadPoolExecutor.CallerRunsPolicy()
        );
        // 允许核心线程超时，节省资源
        COMMON_POOL.allowCoreThreadTimeOut(true);
    }

    /**
     * 出发点（带 TraceId，默认跳过 DFS 检测）
     * @param traceId 链路追踪 ID（用于日志追踪，会自动传递到所有子线程）
     *                如果为 null，会尝试从当前线程的 MDC 获取，如果还没有则自动生成
     * @param timeout 超时时间（毫秒）
     * @param executorService 执行任务的线程池（每次调用独立传入，互不影响）
     * @param workerWrappers 任务包装器列表
     * @return true-所有任务在超时时间内完成，false-发生超时
     */
    public static boolean beginWork(String traceId, long timeout, ExecutorService executorService, List<WorkerWrapper> workerWrappers) throws ExecutionException, InterruptedException {
        return beginWork(traceId, timeout, executorService, false, workerWrappers);
    }
    
    /**
     * 出发点（带 TraceId，可选 DFS 检测）
     * @param traceId 链路追踪 ID（用于日志追踪，会自动传递到所有子线程）
     *                如果为 null，会尝试从当前线程的 MDC 获取，如果还没有则自动生成
     * @param timeout 超时时间（毫秒）
     * @param executorService 执行任务的线程池（每次调用独立传入，互不影响）
     * @param needValidation 是否需要进行 DFS 检测（ID 重复检测 + 循环依赖检测）
     *                       true: 进行检测（安全但有性能损耗）
     *                       false: 跳过检测（高性能，但需要开发人员确保依赖链正确）
     * @param workerWrappers 任务包装器列表
     * @return true-所有任务在超时时间内完成，false-发生超时
     */
    public static boolean beginWork(String traceId, long timeout, ExecutorService executorService, boolean needValidation, List<WorkerWrapper> workerWrappers) throws ExecutionException, InterruptedException {
        // 准备 MDC 上下文（不修改调用者线程的 MDC）
        Map<String, String> mdcContext = prepareMDCContext(traceId);
        
        return doBeginWork(timeout, executorService, needValidation, workerWrappers, mdcContext);
    }
    
    /**
     * 准备 MDC 上下文，不修改调用者线程的 MDC
     * 
     * @param traceId 指定的 TraceId，如果为 null 则从当前线程获取或自动生成
     * @return MDC 上下文（包含 TraceId）
     */
    private static Map<String, String> prepareMDCContext(String traceId) {
        // 获取当前线程的 MDC 上下文（作为基础）
        Map<String, String> mdcContext = TraceContext.getMDCContext();
        
        // 确定要使用的 TraceId
        String finalTraceId;
        if (traceId != null && !traceId.isEmpty()) {
            // 使用传入的 TraceId
            finalTraceId = traceId;
        } else {
            // 从当前线程的 MDC 获取，如果没有则生成新的
            finalTraceId = mdcContext.get(TraceContext.TRACE_ID_KEY);
            if (finalTraceId == null || finalTraceId.isEmpty()) {
                finalTraceId = TraceContext.generateTraceId();
            }
        }
        
        // 将 TraceId 放入 MDC 上下文（不影响调用者线程）
        mdcContext.put(TraceContext.TRACE_ID_KEY, finalTraceId);
        
        return mdcContext;
    }
    
    /**
     * 出发点（不带 TraceId，自动生成，默认跳过 DFS 检测）
     * @param timeout 超时时间（毫秒）
     * @param executorService 执行任务的线程池（每次调用独立传入，互不影响）
     * @param workerWrappers 任务包装器列表
     * @return true-所有任务在超时时间内完成，false-发生超时
     */
    public static boolean beginWork(long timeout, ExecutorService executorService, List<WorkerWrapper> workerWrappers) throws ExecutionException, InterruptedException {
        return beginWork(null, timeout, executorService, workerWrappers);
    }
    
    /**
     * 出发点（不带 TraceId，自动生成，可选 DFS 检测）
     * @param timeout 超时时间（毫秒）
     * @param executorService 执行任务的线程池（每次调用独立传入，互不影响）
     * @param needValidation 是否需要进行 DFS 检测（ID 重复检测 + 循环依赖检测）
     * @param workerWrappers 任务包装器列表
     * @return true-所有任务在超时时间内完成，false-发生超时
     */
    public static boolean beginWork(long timeout, ExecutorService executorService, boolean needValidation, List<WorkerWrapper> workerWrappers) throws ExecutionException, InterruptedException {
        return beginWork(null, timeout, executorService, needValidation, workerWrappers);
    }
    
    /**
     * 实际执行任务的方法
     * @param needValidation 是否需要进行 DFS 检测（ID 重复 + 循环依赖）
     */
    private static boolean doBeginWork(long timeout, ExecutorService executorService, boolean needValidation, List<WorkerWrapper> workerWrappers, Map<String, String> mdcContext) throws ExecutionException, InterruptedException {
        if(workerWrappers == null || workerWrappers.size() == 0) {
            return false;
        }
        
        Set<WorkerWrapper> allWrappers = null;
        
        // 根据 needValidation 参数决定是否进行 DFS 检测
        if (needValidation) {
            // 收集所有 wrapper（包括依赖链）
            allWrappers = new HashSet<>();
            totalWorkers(workerWrappers, allWrappers);
            
            // 检测 ID 重复
            checkDuplicateIds(allWrappers);
            
            // 检测循环依赖
            checkCyclicDependency(allWrappers);
        }
        
        //定义一个map，存放所有的wrapper，key为wrapper的唯一id，value是该wrapper，可以从value中获取wrapper的result
        Map<String, WorkerWrapper> forParamUseWrappers = new ConcurrentHashMap<>();
        CompletableFuture[] futures = new CompletableFuture[workerWrappers.size()];
        for (int i = 0; i < workerWrappers.size(); i++) {
            WorkerWrapper wrapper = workerWrappers.get(i);
            futures[i] = CompletableFuture.runAsync(() -> {
                // 在子线程中设置 MDC 上下文
                TraceContext.setMDCContext(mdcContext);
                try {
                    wrapper.work(executorService, timeout, forParamUseWrappers);
                } finally {
                    // 清理子线程的 MDC，避免线程池复用时污染
                    TraceContext.clear();
                }
            }, executorService);
        }
        try {
            CompletableFuture.allOf(futures).get(timeout, TimeUnit.MILLISECONDS);
            return true;
        } catch (TimeoutException e) {
            Set<WorkerWrapper> set = new HashSet<>();
            totalWorkers(workerWrappers, set);
            for (WorkerWrapper wrapper : set) {
                wrapper.stopNow();
            }
            return false;
        } finally {
            // 清理wrapper的上下文引用，防止内存泄漏
            // 如果之前没有收集 allWrappers（needValidation=false），现在收集用于清理
            if (allWrappers == null) {
                allWrappers = new HashSet<>();
                totalWorkers(workerWrappers, allWrappers);
            }
            for (WorkerWrapper wrapper : allWrappers) {
                wrapper.clearForParamUseWrapper();
            }
        }
    }

    /**
     * 如果想自定义线程池，请传pool（带 TraceId，默认跳过 DFS 检测）
     */
    public static boolean beginWork(String traceId, long timeout, ExecutorService executorService, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        if(workerWrapper == null || workerWrapper.length == 0) {
            return false;
        }
        List<WorkerWrapper> workerWrappers =  Arrays.stream(workerWrapper).collect(Collectors.toList());
        return beginWork(traceId, timeout, executorService, workerWrappers);
    }
    
    /**
     * 如果想自定义线程池，请传pool（带 TraceId，可选 DFS 检测）
     */
    public static boolean beginWork(String traceId, long timeout, ExecutorService executorService, boolean needValidation, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        if(workerWrapper == null || workerWrapper.length == 0) {
            return false;
        }
        List<WorkerWrapper> workerWrappers =  Arrays.stream(workerWrapper).collect(Collectors.toList());
        return beginWork(traceId, timeout, executorService, needValidation, workerWrappers);
    }
    
    /**
     * 如果想自定义线程池，请传pool。不自定义的话，就走默认的COMMON_POOL（默认跳过 DFS 检测）
     */
    public static boolean beginWork(long timeout, ExecutorService executorService, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        return beginWork(null, timeout, executorService, workerWrapper);
    }
    
    /**
     * 如果想自定义线程池，请传pool。不自定义的话，就走默认的COMMON_POOL（可选 DFS 检测）
     */
    public static boolean beginWork(long timeout, ExecutorService executorService, boolean needValidation, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        return beginWork(null, timeout, executorService, needValidation, workerWrapper);
    }

    /**
     * 同步阻塞,直到所有都完成,或失败（带 TraceId，默认跳过 DFS 检测）
     */
    public static boolean beginWork(String traceId, long timeout, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        return beginWork(traceId, timeout, COMMON_POOL, workerWrapper);
    }
    
    /**
     * 同步阻塞,直到所有都完成,或失败（带 TraceId，可选 DFS 检测）
     */
    public static boolean beginWork(String traceId, long timeout, boolean needValidation, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        return beginWork(traceId, timeout, COMMON_POOL, needValidation, workerWrapper);
    }

    /**
     * 同步阻塞,直到所有都完成,或失败（默认跳过 DFS 检测）
     */
    public static boolean beginWork(long timeout, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        return beginWork(null, timeout, COMMON_POOL, workerWrapper);
    }
    
    /**
     * 同步阻塞,直到所有都完成,或失败（可选 DFS 检测）
     */
    public static boolean beginWork(long timeout, boolean needValidation, WorkerWrapper... workerWrapper) throws ExecutionException, InterruptedException {
        return beginWork(null, timeout, COMMON_POOL, needValidation, workerWrapper);
    }

    public static void beginWorkAsync(String traceId, long timeout, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        beginWorkAsync(traceId, timeout, COMMON_POOL, groupCallback, workerWrapper);
    }
    
    public static void beginWorkAsync(String traceId, long timeout, boolean needValidation, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        beginWorkAsync(traceId, timeout, COMMON_POOL, needValidation, groupCallback, workerWrapper);
    }

    public static void beginWorkAsync(long timeout, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        beginWorkAsync(null, timeout, COMMON_POOL, groupCallback, workerWrapper);
    }
    
    public static void beginWorkAsync(long timeout, boolean needValidation, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        beginWorkAsync(null, timeout, COMMON_POOL, needValidation, groupCallback, workerWrapper);
    }

    /**
     * 异步执行,直到所有都完成,或失败后，发起回调（带 TraceId，默认跳过 DFS 检测）
     */
    public static void beginWorkAsync(String traceId, long timeout, ExecutorService executorService, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        beginWorkAsync(traceId, timeout, executorService, false, groupCallback, workerWrapper);
    }
    
    /**
     * 异步执行,直到所有都完成,或失败后，发起回调（带 TraceId，可选 DFS 检测）
     */
    public static void beginWorkAsync(String traceId, long timeout, ExecutorService executorService, boolean needValidation, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        if (groupCallback == null) {
            groupCallback = new DefaultGroupCallback();
        }
        IGroupCallback finalGroupCallback = groupCallback;
        
        // 准备 MDC 上下文（不修改调用者线程的 MDC）
        Map<String, String> mdcContext = prepareMDCContext(traceId);
        
        if (executorService != null) {
            executorService.submit(() -> {
                // 在子线程中设置 MDC
                TraceContext.setMDCContext(mdcContext);
                try {
                    boolean success = beginWork(traceId, timeout, executorService, needValidation, workerWrapper);
                    if (success) {
                        finalGroupCallback.success(Arrays.asList(workerWrapper));
                    } else {
                        finalGroupCallback.failure(Arrays.asList(workerWrapper), new TimeoutException());
                    }
                } catch (ExecutionException | InterruptedException e) {
                    logger.error("异步任务执行失败", e);
                    finalGroupCallback.failure(Arrays.asList(workerWrapper), e);
                } finally {
                    TraceContext.clear();
                }
            });
        } else {
            COMMON_POOL.submit(() -> {
                // 在子线程中设置 MDC
                TraceContext.setMDCContext(mdcContext);
                try {
                    boolean success = beginWork(traceId, timeout, COMMON_POOL, needValidation, workerWrapper);
                    if (success) {
                        finalGroupCallback.success(Arrays.asList(workerWrapper));
                    } else {
                        finalGroupCallback.failure(Arrays.asList(workerWrapper), new TimeoutException());
                    }
                } catch (ExecutionException | InterruptedException e) {
                    logger.error("异步任务执行失败", e);
                    finalGroupCallback.failure(Arrays.asList(workerWrapper), e);
                } finally {
                    TraceContext.clear();
                }
            });
        }
    }
    
    /**
     * 异步执行,直到所有都完成,或失败后，发起回调（默认跳过 DFS 检测）
     */
    public static void beginWorkAsync(long timeout, ExecutorService executorService, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        beginWorkAsync(null, timeout, executorService, groupCallback, workerWrapper);
    }
    
    /**
     * 异步执行,直到所有都完成,或失败后，发起回调（可选 DFS 检测）
     */
    public static void beginWorkAsync(long timeout, ExecutorService executorService, boolean needValidation, IGroupCallback groupCallback, WorkerWrapper... workerWrapper) {
        beginWorkAsync(null, timeout, executorService, needValidation, groupCallback, workerWrapper);
    }

    /**
     * 总共多少个执行单元
     */
    @SuppressWarnings("unchecked")
    private static void totalWorkers(List<WorkerWrapper> workerWrappers, Set<WorkerWrapper> set) {
        set.addAll(workerWrappers);
        for (WorkerWrapper wrapper : workerWrappers) {
            if (wrapper.getNextWrappers() == null) {
                continue;
            }
            List<WorkerWrapper> wrappers = wrapper.getNextWrappers();
            totalWorkers(wrappers, set);
        }
    }
    
    /**
     * 检测 Wrapper ID 是否重复
     * @param allWrappers 所有的 wrapper
     * @throws IllegalArgumentException 如果发现重复的 ID
     */
    private static void checkDuplicateIds(Set<WorkerWrapper> allWrappers) {
        Map<String, WorkerWrapper> idMap = new HashMap<>();
        List<String> duplicateIds = new ArrayList<>();
        
        for (WorkerWrapper wrapper : allWrappers) {
            String id = wrapper.getId();
            if (idMap.containsKey(id)) {
                duplicateIds.add(id);
                logger.error("发现重复的 Wrapper ID: {}, 这会导致结果混乱！", id);
            } else {
                idMap.put(id, wrapper);
            }
        }
        
        if (!duplicateIds.isEmpty()) {
            String errorMsg = "存在重复的 Wrapper ID: " + duplicateIds + 
                            ". 请确保每个 WorkerWrapper 的 ID 唯一，或者不手动指定 ID（使用默认的 UUID）。";
            logger.error(errorMsg);
            throw new IllegalArgumentException(errorMsg);
        }
        
        logger.debug("ID 检测通过，共 {} 个 wrapper，无重复", allWrappers.size());
    }
    
    /**
     * 检测循环依赖
     * 在 beginWork 时统一检测整个 workflow 图，避免在 build 时重复检测
     * 
     * @param allWrappers 所有的 wrapper 集合
     * @throws IllegalStateException 如果发现循环依赖
     */
    private static void checkCyclicDependency(Set<WorkerWrapper> allWrappers) {
        Set<WorkerWrapper> visited = new HashSet<>();
        Set<WorkerWrapper> recursionStack = new HashSet<>();
        
        // 对每个 wrapper 进行 DFS 检测
        for (WorkerWrapper wrapper : allWrappers) {
            if (!visited.contains(wrapper)) {
                List<WorkerWrapper> path = new ArrayList<>();
                if (hasCycle(wrapper, visited, recursionStack, path)) {
                    // 构建循环路径信息
                    StringBuilder cycleInfo = new StringBuilder("检测到循环依赖: ");
                    for (int i = 0; i < path.size(); i++) {
                        cycleInfo.append(path.get(i).getId());
                        if (i < path.size() - 1) {
                            cycleInfo.append(" -> ");
                        }
                    }
                    logger.error(cycleInfo.toString());
                    throw new IllegalStateException(cycleInfo.toString());
                }
            }
        }
        
        logger.debug("循环依赖检测通过，共 {} 个 wrapper，无循环", allWrappers.size());
    }
    
    /**
     * DFS 检测是否存在环
     * 
     * @param current 当前节点
     * @param visited 已访问的节点集合
     * @param recursionStack 当前递归栈（用于检测环）
     * @param path 当前路径（用于记录循环路径）
     * @return true 如果存在环，false 如果不存在环
     */
    private static boolean hasCycle(WorkerWrapper current, Set<WorkerWrapper> visited, 
                                    Set<WorkerWrapper> recursionStack, List<WorkerWrapper> path) {
        if (current == null) {
            return false;
        }
        
        // 如果当前节点在递归栈中，说明找到了环
        if (recursionStack.contains(current)) {
            // 截取环的部分
            int cycleStart = path.indexOf(current);
            if (cycleStart >= 0) {
                path.subList(0, cycleStart).clear();
            }
            path.add(current);
            return true;
        }
        
        // 如果已经访问过且不在递归栈中，说明这个节点已经检测过了，没有环
        if (visited.contains(current)) {
            return false;
        }
        
        // 标记为已访问并加入递归栈
        visited.add(current);
        recursionStack.add(current);
        path.add(current);
        
        // 只检查下游节点（沿执行方向）
        // 由于外层循环已遍历所有节点，只需单向检查即可检测所有环
        // 如果同时检查上游和下游，会把有向图当作无向图，导致误报
        if (current.getNextWrappers() != null) {
            for (Object obj : current.getNextWrappers()) {
                WorkerWrapper<?, ?> nextWrapper = (WorkerWrapper<?, ?>) obj;
                if (hasCycle(nextWrapper, visited, recursionStack, path)) {
                    return true;
                }
            }
        }
        
        // 从递归栈中移除（回溯）
        recursionStack.remove(current);
        path.remove(path.size() - 1);
        
        return false;
    }

    /**
     * 关闭默认线程池（COMMON_POOL）
     * 注意：如果使用自定义线程池，请自行管理其生命周期
     */
    public static void shutDown() {
        shutDown(COMMON_POOL);
    }

    /**
     * 优雅关闭指定的线程池
     * @param executorService 要关闭的线程池
     */
    public static void shutDown(ExecutorService executorService) {
        shutDown(executorService, 60, TimeUnit.SECONDS);
    }
    
    /**
     * 优雅关闭线程池，带超时控制
     * @param executorService 要关闭的线程池
     * @param timeout 等待超时时间
     * @param unit 时间单位
     */
    public static void shutDown(ExecutorService executorService, long timeout, TimeUnit unit) {
        if (executorService == null) {
            return;
        }
        
        executorService.shutdown();
        try {
            // 等待已提交的任务完成
            if (!executorService.awaitTermination(timeout, unit)) {
                logger.warn("线程池在{}{}内未能正常关闭，尝试强制关闭", timeout, unit);
                // 超时后强制关闭
                executorService.shutdownNow();
                // 再等待一小段时间
                if (!executorService.awaitTermination(10, TimeUnit.SECONDS)) {
                    logger.error("线程池强制关闭失败，可能存在无法中断的任务");
                }
            } else {
                logger.info("线程池已优雅关闭");
            }
        } catch (InterruptedException e) {
            logger.warn("线程池关闭过程被中断，立即强制关闭", e);
            // 中断时立即强制关闭
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    public static String getThreadCount() {
        return "activeCount=" + COMMON_POOL.getActiveCount() +
                "  completedCount " + COMMON_POOL.getCompletedTaskCount() +
                "  largestCount " + COMMON_POOL.getLargestPoolSize();
    }
}
